package client

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"zenoLink/conf"
	"zenoLink/protocol"
	"zenoLink/protocol/pb"
	"zenoLink/utils"

	"google.golang.org/protobuf/proto"
)

type socket5Client struct {
	RemoteHost string
	RemotePort uint16
	SessionId  string
}

func NewSocket5Client(remoteHost string, remotePort uint16) InterClient {
	client := &socket5Client{}
	client.RemoteHost = remoteHost
	client.RemotePort = remotePort
	return client
}

func (m *socket5Client) Process(client net.Conn) {
	if err := m.Socks5Auth(client); err != nil {
		slog.Debug("socket auth err", err.Error())
		client.Close()
		return
	}

	target, err := m.Socks5Connect(client)
	if err != nil {
		fmt.Println("Socks5Connect err=", err)
		client.Close()
		return
	}

	m.Socks5Forward(client, target)
}

func (m *socket5Client) Socks5Auth(client net.Conn) (err error) {
	buf := make([]byte, 256)

	// 读取 VER 和 NMETHODS
	n, err := io.ReadFull(client, buf[:2])
	if n != 2 {
		return errors.New("reading header: " + err.Error())
	}

	ver, nMethods := int(buf[0]), int(buf[1])
	if ver != 5 {
		return errors.New("invalid version")
	}

	// 读取 METHODS 列表
	n, err = io.ReadFull(client, buf[:nMethods])
	if n != nMethods {
		return errors.New("reading methods: " + err.Error())
	}

	//无需认证
	n, err = client.Write([]byte{0x05, 0x00})
	if n != 2 || err != nil {
		return errors.New("write rsp err: " + err.Error())
	}

	return nil
}

func (m *socket5Client) Socks5Connect(client net.Conn) (net.Conn, error) {
	buf := make([]byte, 256)

	n, err := io.ReadFull(client, buf[:4])
	if n != 4 {
		return nil, errors.New("read header: " + err.Error())
	}

	ver, cmd, _, atyp := buf[0], buf[1], buf[2], buf[3]
	if ver != 5 || cmd != 1 {
		return nil, errors.New("invalid ver/cmd")
	}

	//TO BE CONTINUED...

	addr := ""
	switch atyp {
	case 1:
		n, err = io.ReadFull(client, buf[:4])
		if n != 4 {
			return nil, errors.New("invalid IPv4: " + err.Error())
		}
		addr = fmt.Sprintf("%d.%d.%d.%d", buf[0], buf[1], buf[2], buf[3])

	case 3:
		n, err = io.ReadFull(client, buf[:1])
		if n != 1 {
			return nil, errors.New("invalid hostname: " + err.Error())
		}
		addrLen := int(buf[0])

		n, err = io.ReadFull(client, buf[:addrLen])
		if n != addrLen {
			return nil, errors.New("invalid hostname: " + err.Error())
		}
		addr = string(buf[:addrLen])

	case 4:
		return nil, errors.New("IPv6: no supported yet")

	default:
		return nil, errors.New("invalid atyp")
	}

	n, err = io.ReadFull(client, buf[:2])
	if n != 2 {
		return nil, errors.New("read port: " + err.Error())
	}
	port := binary.BigEndian.Uint16(buf[:2])

	dest, err := m.RemoteConnect(addr, port)
	if err != nil {
		return nil, err
	}

	_, err = client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
	if err != nil {
		dest.Close()
		return nil, errors.New("write rsp: " + err.Error())
	}

	return dest, nil
}

func (m *socket5Client) RemoteConnect(reqHost string, reqPort uint16) (net.Conn, error) {

	mesh := pb.OutMeshConnet{}
	mesh.SessionId = m.SessionId
	mesh.Host = reqHost
	mesh.Port = int32(reqPort)

	destAddrPort := fmt.Sprintf("%s:%d", m.RemoteHost, m.RemotePort)
	dest, err := net.Dial("tcp", destAddrPort)
	if err != nil {
		return nil, errors.New("dial dst: " + err.Error())
	}

	bytes, err := proto.Marshal(&mesh)
	if err != nil {
		dest.Close()
		return nil, err
	}

	enBytes, err := utils.Encrypt(bytes, conf.EncryptKey)
	if err != nil {
		return nil, err
	}

	data := utils.NetOutPackUint16(enBytes)
	_, err = dest.Write(data)
	if err != nil {
		dest.Close()
		return nil, err
	}

	data = make([]byte, 2)
	n, err := dest.Read(data)
	if err != nil {
		return nil, err
	}
	if n != 2 {
		err = fmt.Errorf("read data is not 2")
		return nil, err
	}

	len := binary.BigEndian.Uint16(data)
	if len > conf.MaxLenPackage {
		err = fmt.Errorf("connection data length is too big")
		return nil, err
	}

	data = make([]byte, len)
	_, err = dest.Read(data)
	if err != nil {
		return nil, err
	}

	deData, err := utils.Decrypt(data, conf.EncryptKey)
	if err != nil {
		return nil, err
	}
	res := &pb.OutMeshConnetRes{}
	err = proto.Unmarshal(deData, res)
	if err != nil {
		return nil, err
	}

	if res.Code != protocol.CodeOk {
		err = fmt.Errorf("code is %v not ok", res.Code)
		return nil, err

	}

	return dest, nil
}

func (m *socket5Client) Socks5Forward(client, target net.Conn) {

	go m.Raw2Encode(client, target)
	go m.Encode2Raw(target, client)
}

func (m *socket5Client) Raw2Encode(source, targert net.Conn) {

	defer targert.Close()
	defer source.Close()

	// reader := bufio.NewReader(source)

	buf := make([]byte, 1300)
	for {
		n, err := source.Read(buf)
		if err != nil {
			if err != io.EOF {
				slog.Error("err=", err.Error(), n)
			}

			break
		}

		trans := pb.OutMeshTransfer{}
		trans.SessionId = m.SessionId
		trans.Data = buf[:n]

		bytes, err := proto.Marshal(&trans)
		if err != nil {
			slog.Error(fmt.Sprintf("proto marshal err %v", err))
			break
		}

		enbytes, err := utils.Encrypt(bytes, conf.EncryptKey)
		if err != nil {
			slog.Error(fmt.Sprintf("utils.Encrypt err=%v", err))
			break
		}

		outData := utils.NetOutPackUint16(enbytes)
		targert.Write(outData)
	}

}

func (m *socket5Client) Encode2Raw(source, targert net.Conn) {
	defer targert.Close()
	defer source.Close()

	// reader := bufio.NewReader(source)

	buf := make([]byte, 2)
	var n int
	var err error
	for {
		n, err = source.Read(buf)
		if err != nil {
			if err != io.EOF {
				slog.Error("err=", err.Error(), n)
			}

			break
		}
		if n != 2 {
			slog.Error(fmt.Sprintf("read len n is not 2 n=%v", n))
			break
		}

		len := binary.BigEndian.Uint16(buf[:2])
		if len > conf.MaxLenPackage {
			slog.Error("data is too big")
		}
		data := make([]byte, len)
		n, err = source.Read(data)
		if err != nil {
			break
		}

		deData, err := utils.Decrypt(data[:n], conf.EncryptKey)
		if err != nil {
			break
		}

		trans := pb.OutMeshTransfer{}
		err = proto.Unmarshal(deData, &trans)
		if err != nil {
			slog.Error(fmt.Sprintf("proto marshal err %v", err))
			break
		}

		targert.Write(trans.Data)
	}

}
